// Copyright (C) 2024 Kumo inc.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <ktest/ktest.h>
#include <atomic>
#include <turbo/times/time.h>
#include <turbo/base/macros.h>
#include <turbo/log/logging.h>
#include <turbo/threading/thread_local.h>
#include <kthread/internal/kutex.h>
#include <turbo/log/logging.h>
#include <kthread/internal/kthread.h>
#include <kthread/internal/task_control.h>
#include <turbo/flags/reflection.h>
#include <kthread/internal/config.h>

namespace kthread {
    extern TaskControl *g_task_control;
}

namespace {
    void *dummy(void *) {
        return nullptr;
    }

    TEST(KthreadTest, setconcurrency) {
        ASSERT_EQ(8 + KTHREAD_EPOLL_THREAD_NUM + KTHREAD_TASK_EXECUTOR_THREAD_NUM, (size_t) kthread_getconcurrency());
        ASSERT_EQ(EINVAL, kthread_setconcurrency(KTHREAD_MIN_CONCURRENCY - 1));
        ASSERT_EQ(EINVAL, kthread_setconcurrency(0));
        ASSERT_EQ(EINVAL, kthread_setconcurrency(-1));
        ASSERT_EQ(EINVAL, kthread_setconcurrency(KTHREAD_MAX_CONCURRENCY + 1));
        ASSERT_EQ(0, kthread_setconcurrency(KTHREAD_MIN_CONCURRENCY));
        ASSERT_EQ(KTHREAD_MIN_CONCURRENCY, kthread_getconcurrency());
        ASSERT_EQ(0, kthread_setconcurrency(KTHREAD_MIN_CONCURRENCY + 1));
        ASSERT_EQ(KTHREAD_MIN_CONCURRENCY + 1, kthread_getconcurrency());
        ASSERT_EQ(0, kthread_setconcurrency(KTHREAD_MIN_CONCURRENCY));  // smaller value
        kthread_t th;
        ASSERT_EQ(0, kthread_start_urgent(&th, nullptr, dummy, nullptr));
        ASSERT_EQ(KTHREAD_MIN_CONCURRENCY + 1, kthread_getconcurrency());
        ASSERT_EQ(0, kthread_setconcurrency(KTHREAD_MIN_CONCURRENCY + 5));
        ASSERT_EQ(KTHREAD_MIN_CONCURRENCY + 5, kthread_getconcurrency());
        ASSERT_EQ(EPERM, kthread_setconcurrency(KTHREAD_MIN_CONCURRENCY + 1));
        ASSERT_EQ(KTHREAD_MIN_CONCURRENCY + 5, kthread_getconcurrency());
    }

    static std::atomic<int> *odd;
    static std::atomic<int> *even;

    static std::atomic<int> nkthreads(0);
    static std::atomic<int> npthreads(0);
    static TURBO_THREAD_LOCAL bool counted = false;
    static std::atomic<bool> stop(false);

    static void *odd_thread(void *) {
        nkthreads.fetch_add(1);
        while (!stop) {
            if (!counted) {
                counted = true;
                npthreads.fetch_add(1);
            }
            kthread::kutex_wake_all(even);
            kthread::kutex_wait(odd, 0, nullptr);
        }
        return nullptr;
    }

    static void *even_thread(void *) {
        nkthreads.fetch_add(1);
        while (!stop) {
            if (!counted) {
                counted = true;
                npthreads.fetch_add(1);
            }
            kthread::kutex_wake_all(odd);
            kthread::kutex_wait(even, 0, nullptr);
        }
        return nullptr;
    }

    TEST(KthreadTest, setconcurrency_with_running_kthread) {
        odd = kthread::kutex_create_checked<std::atomic<int> >();
        even = kthread::kutex_create_checked<std::atomic<int> >();
        ASSERT_TRUE(odd != nullptr && even != nullptr);
        *odd = 0;
        *even = 0;
        std::vector<kthread_t> tids;
        const int N = 500;
        for (int i = 0; i < N; ++i) {
            kthread_t tid;
            kthread_start_background(&tid, &KTHREAD_ATTR_SMALL, odd_thread, nullptr);
            tids.push_back(tid);
            kthread_start_background(&tid, &KTHREAD_ATTR_SMALL, even_thread, nullptr);
            tids.push_back(tid);
        }
        for (int i = 100; i <= N; ++i) {
            ASSERT_EQ(0, kthread_setconcurrency(i));
            ASSERT_EQ(i, kthread_getconcurrency());
        }
        usleep(1000 * N);
        *odd = 1;
        *even = 1;
        stop = true;
        kthread::kutex_wake_all(odd);
        kthread::kutex_wake_all(even);
        for (size_t i = 0; i < tids.size(); ++i) {
            kthread_join(tids[i], nullptr);
        }
        KLOG(INFO) << "All kthreads has quit";
        ASSERT_EQ(2 * N, nkthreads);
        // This is not necessarily true, not all workers need to run sth.
        //ASSERT_EQ(N, npthreads);
        KLOG(INFO) << "Touched pthreads=" << npthreads;
    }

    void *sleep_proc(void *) {
        usleep(100000);
        return nullptr;
    }

    void *add_concurrency_proc(void *) {
        kthread_t tid;
        kthread_start_background(&tid, &KTHREAD_ATTR_SMALL, sleep_proc, nullptr);
        kthread_join(tid, nullptr);
        return nullptr;
    }

    bool set_min_concurrency(int num) {
        std::stringstream ss;
        ss << num;
        auto info = turbo::find_command_line_flag("kthread_min_concurrency");
        if(!info) {
            return false;
        }
        std::string err;
        auto r = info->user_validate(ss.str(), &err);
        if(!r) {
            return false;
        }
        turbo::set_flag(&FLAGS_kthread_min_concurrency, num);
        return true;
    }

    int get_min_concurrency() {
        return  turbo::get_flag(FLAGS_kthread_min_concurrency);
    }

    TEST(KthreadTest, min_concurrency) {
        KLOG(WARNING)<<"min_concurrency------------------------";
        ASSERT_EQ(1, set_min_concurrency(-1)); // set min success
        ASSERT_EQ(1, set_min_concurrency(0)); // set min success
        ASSERT_EQ(0, get_min_concurrency());
        int conn = kthread_getconcurrency();
        int add_conn = 100;
        ASSERT_EQ(0, set_min_concurrency(conn + 1)); // set min failed
        ASSERT_EQ(1, set_min_concurrency(conn - 1)); // set min success
        ASSERT_EQ(conn - 1, get_min_concurrency());
        KLOG(WARNING)<<"min_concurrency------------------------";
        ASSERT_EQ(EINVAL, kthread_setconcurrency(conn - 2)); // set max failed
        ASSERT_EQ(0, kthread_setconcurrency(conn + add_conn + 1)); // set max success
        KLOG(WARNING)<<"min_concurrency######";
        ASSERT_EQ(0, kthread_setconcurrency(conn + add_conn)); // set max success
        ASSERT_EQ(conn + add_conn, kthread_getconcurrency());
        ASSERT_EQ(conn, kthread::g_task_control->concurrency());
        auto old = get_min_concurrency();
        set_min_concurrency(conn + 1);
        ASSERT_NE(old, get_min_concurrency()); // set min success
        ASSERT_EQ(conn + 1, get_min_concurrency());
        ASSERT_EQ(conn + 1, kthread::g_task_control->concurrency());
        KLOG(WARNING)<<"min_concurrency------------------------";
        std::vector<kthread_t> tids;
        for (int i = 0; i < conn; ++i) {
            kthread_t tid;
            kthread_start_background(&tid, &KTHREAD_ATTR_SMALL, sleep_proc, nullptr);
            tids.push_back(tid);
        }
        for (int i = 0; i < add_conn; ++i) {
            kthread_t tid;
            kthread_start_background(&tid, &KTHREAD_ATTR_SMALL, add_concurrency_proc, nullptr);
            tids.push_back(tid);
        }
        for (size_t i = 0; i < tids.size(); ++i) {
            kthread_join(tids[i], nullptr);
        }
        ASSERT_EQ(conn + add_conn, kthread_getconcurrency());
        ASSERT_EQ(conn + add_conn, kthread::g_task_control->concurrency());
    }

    int current_tag(int tag) {
        std::stringstream ss;
        ss << tag;
        auto info = turbo::find_command_line_flag("kthread_current_tag");
        if(!info) {
            return false;
        }
        std::string err;
        auto r = info->user_validate(ss.str(), &err);
        if(!r) {
            return false;
        }
        turbo::set_flag(&FLAGS_kthread_current_tag, tag);
        return true;
    }

    TEST(KthreadTest, current_tag) {
        ASSERT_EQ(false, current_tag(-1));
        ASSERT_EQ(true, current_tag(0));
        ASSERT_EQ(false, current_tag(1));
    }

    int concurrency_by_tag(int num) {
        std::stringstream ss;
        ss << num;
        auto info = turbo::find_command_line_flag("kthread_concurrency_by_tag");
        if(!info) {
            return false;
        }
        std::string err;
        auto r = info->user_validate(ss.str(), &err);
        if(!r) {
            return false;
        }
        turbo::set_flag(&FLAGS_kthread_concurrency_by_tag, num);
        return true;
    }

    TEST(KthreadTest, concurrency_by_tag) {
        ASSERT_EQ(concurrency_by_tag(1), true);
        ASSERT_EQ(concurrency_by_tag(1), false);
        auto con = kthread_getconcurrency_by_tag(0);
        ASSERT_EQ(concurrency_by_tag(con), true);
        ASSERT_EQ(concurrency_by_tag(con + 1), false);
        kthread_setconcurrency(con + 1);
        ASSERT_EQ(concurrency_by_tag(con + 1), true);
    }

} // namespace
